#---------------------安装所需的pip包--------------------------#
# pip install cnocr -i https://pypi.doubanio.com/simple        #
# pip install PyMuPDF                                          #
#--------------------------------------------------------------#

#!/usr/bin/python
# -*- coding: utf-8 -*-

import fitz
from cnocr.utils import read_img
from cnocr import CnOcr
import os
import sys
import shutil
import ctypes

#---把pdf每页都转换成一张图片---#
def pdf_to_img(input_pdf_name, output_file):
    pdf = fitz.open(input_pdf_name + '.pdf')
    page = pdf.page_count #读取pdf的页数
    for i in range (0, page):
        page_fn = pdf.load_page(i) #读取第i页
        pix = page_fn.get_pixmap()
        pix.save(output_file + '/' + str(i) + '.png') #从0开始按序号另存页数为图片

        
#---识别图片中的文字并且return---#
def img_ocr(input_pic):
    ocr = CnOcr()
    img_fp = input_pic #输入图片路径
    img = read_img(img_fp) #读取图片
    res = ocr.ocr(img) #res是list（列表）
    len_res = len(res) #文本的行数
    text = ""
    for i in range (0,len_res):
        line_tuple = res[i] #line_tuple是tuple（元组）
        len_line_tuple = len(line_tuple[0])
        s = ""
        single_line = s.join(line_tuple[0])
        text += single_line
    return text #text为识别出来的文本，此处return可以给下面写入txt

#---将识别出来的文字写入txt---#
def write_to_txt(input_pdf_name,text):
    #file = open(input_pdf_name + '.txt','a',encoding='utf-8') #用a模式是续写而不是w(重新写入),编码为utf-8
    with open (input_pdf_name + '.txt', 'a', encoding='utf-8') as file:    
        file.write(text)
        file.write('\n') #写完一张图片里面的内容就换行   
        
#---获取当前路径下所有pdf文件---#        
def get_pdf_list():
    pdf_list = [x for x in os.listdir('.') if os.path.isfile(x) and os.path.splitext(x)[1]=='.pdf']
    return pdf_list

pdf_list = get_pdf_list() #获取当前路径下所有pdf的文件，返回值是list（列表） ##剩下代码均归入pdf_list的循环中（v2.0增加此功能）
for i in range(0,len(pdf_list)):
    name = os.path.splitext(pdf_list[i])[0]

    input_pdf_name = './' + name
    output_file = './pdf_2_image' #临时生成的文件夹的名字（可以随意更改）
    input_file = output_file #pdf_to_img的output_file == img_ocr的input_pic
    os.mkdir(output_file) #生成临时文件夹存放pdf另存出来的img

    pdf_to_img(input_pdf_name, output_file)
    pic = os.listdir(input_file) #遍历临时文件夹下的图片,pic是list（列表）
    for pic_fn in pic: 
        text = img_ocr(input_file + '/' + pic_fn)
        write_to_txt(input_pdf_name,text)
    shutil.rmtree(output_file) #删除生成的临时文件夹和文件
    #print('已完成 {}.pdf 的转换'.format(name)) #v3.0注释掉此行
    #print('已将 {}.pdf 转化为 {}.txt'.format(name,name)) # v3.0增加此行
#print('全部转换均已完成')
ctypes.windll.user32.MessageBoxW(0, u'已完成PDF转化', u'已完成',0)